package com.thghh.ch5.huffman;

import java.util.*;

/**
 * @author Zhikang.Peng
 * @version 1.0
 * @email thghh@qq.com
 * @date 2020/9/28 16:21
 */
public class HuffmanTree<E> {

    private Node<E> root;

    public HuffmanTree(Map<E, Double> map) {
        generateTree(map);
    }

    public Node<E> getRoot() {
        return root;
    }

    private void generateTree(Map<E, Double> map) {
        Stack<Node> nodeList = new Stack<>();
        for(Map.Entry<E, Double> entry : map.entrySet()) {
            Node<E> node = new Node<>(entry.getKey(), entry.getValue());
            nodeList.add(node);
        }
        while(nodeList.size() > 1) {
            sort(nodeList);
            Node<E> left = nodeList.pop();
            Node<E> right = nodeList.pop();
            String key = String.format("%s%s", left.data,  right.data);
            double d = left.weight + right.weight;
            Node parent = new Node(key, d);

            parent.left = left;
            parent.right = right;
            left.parent = parent;
            right.parent = parent;

            root = parent;

            nodeList.add(parent);
        }
    }

    private void sort(List<Node> nodeList) {
        Collections.sort(nodeList, (t1,t2) -> {
            if(t1.weight < t2.weight) return 1;
            if(t1.weight > t2.weight) return -1;
            return 0;
        });
    }

    public static class Node<E> implements Traverse<E>{
        E data;
        double weight;
        Node<E> left;
        Node<E> right;
        Node<E> parent;

        public Node(E data, double weight) {
            this.data = data;
            this.weight = weight;
        }
        @Override
        public Node<E> getLeft() {
            return left;
        }

        @Override
        public Node<E> getRight() {
            return right;
        }

        @Override
        public E getValue() {
            return data;
        }

        @Override
        public Node<E> getObject() {
            return this;
        }

        public String toString() {
            return "Node[data=" + data + ", weight=" + weight + "]";
        }
    }
}
